'''
Created on Feb 2, 2010

@author: oabalbin
'''

#import sys;
#import string;
#import StringIO;
import os;
import sys;
import glob;
import MySQLdb;
import numpy as np
from optparse import OptionParser
from collections import deque

# My libraries
import signatures.parsers.read_microarray_data as pm
import signatures.parsers.read_gene_lists as gpm
import signatures.db.tables as dt
import signatures.db.query as dq
import signatures.common.classes as cm #record, arrayMatrix
import signatures.preprocess.tools as tl #filter_constant_genes, get_variable_genes, 
import signatures.seedgenes.seedlist as sd #get_seed_genes, write_varin_file
import signatures.bfrm.paramfile as bp
import signatures.bfrm.bfrm_run as br

class pipeline():
    
    def __init__(self, dbhost, dbuser, dbpasswd, database):
        
        self.mt = dt.tables(dbhost, dbuser, dbpasswd, database)
        self.mq = dq.query(dbhost, dbuser, dbpasswd, database)
        #table information
        self.tableDescription = 'probeID VARCHAR(40), geneSym VARCHAR(40), sample VARCHAR(40), log_fold FLOAT(20,14), INDEX probe (probeID), INDEX gene (geneSym), INDEX sample (sample)'
        self.tableName = 'ERG_Prostate'
        #self.tableName = 'erg_signatures_only'
        # Parameters for pre-processing data:
        # percentage of low variable gene to be filter out of the analysis
        self.lvar=25
        # percentage of low median gene to be filter out of the analysis
        self.lmed=25
        
    
    
    def parse_input_annotfiles(self, annotInputFile):
    
        expr_file, seed_file, outdirs, db_file, samples, sampleclass, typefile={},{},{},{},{},{},'nofile'
    
        for line in annotInputFile:
            fields = line.strip('\n')
                 
            if fields[0][0] == '':
                continue
            
            if fields[0][0] == '#':
                typefile = fields
                continue
            
            fields = fields.split('=')
            dict = {'#Array expression':expr_file, '#Seed genes':seed_file,'#Database file':db_file,'#Samples':samples,'#Output Directories':outdirs,\
                    '#Sample classes':sampleclass}[typefile]
            if typefile ==  '#Sample classes':
                spclass = fields[1].split('\t')
                dict[fields[0]] = spclass
            else:
                dict[fields[0]] = fields[1]
            
            
        return expr_file, seed_file, db_file, samples, sampleclass, outdirs
    

    def read_files_folder(self,folderpath,ext):
        ''' '''
        # Read files in folder
        myfiles=[]
        for infile in glob.glob( os.path.join(folderpath, '*'+ext) ):
            myfiles.append(infile)
            #print "current file is: " + infile
        return myfiles
    
    
    def process_microarray_files(self,exp_files,tp):
        dbf=[]
        for f in exp_files:
            of=f+'_dbfile'
            tp.microarray_exp_file(open(f),open(of,'w'))
            dbf.append(of)
        
        return dbf
    
    def create_database(self,dbfiles,doit=False):
        """
        It creates a databases if doit = True
        """
        if doit:
            print 'Making tables'            
            for dbf in dbfiles:     
                if not(self.mt.table_existence(self.tableName)):
                    self.mt.create_table(self.tableName,self.tableDescription)
                    
                self.mt.load_data_table(self.tableName,dbf)
    
    def initialize_database(self, exp_files,tp, doit=False):
        
        dbfiles = self.process_microarray_files(exp_files,tp)
        self.create_database(dbfiles,doit)
            
    
    def filter_expMat(self, samplelist, initgenelist, tableName, seedfile, dofilter=False):     
        '''
        Returns an arrayMatrix with the final expression matrix.
        From the matrix you: 
        
        a) filter the all genes to remove noise, keep only genes var >= lvar, med >= lmed
        b) if another optional list is provided, you want to see
           b.1) how many of the new genes are variable. 
           b.2) Produce a expression matrix with only genes in b.1
        
        '''
        self.mt.create_connector()
        self.mq.create_connector()
        # Expression array matrix.            
        expMat = self.mq.get_exp_mat_2(samplelist, initgenelist, self.tableName)
        # Normalize the expression values data, using oncomine approach.
        expMat.normalization()
        
        if dofilter:
            # Find variable genes and modify he expMat.
            genelist = tl.filter_out_constant_genes(expMat, self.lvar, self.lmed, False)       
            # Get the index and name of the variable genes for a particular experiment
            mgenelist, dictmgl =  tl.get_variable_genes(expMat,genelist)
            
            expMat = tl.get_filtered_expMat(expMat,dictmgl,len(samplelist))
            expMat.create_gen_index()
     
        # list of external genes to be used as the "universe of genes"
        # e.i: this could be the ERG interactome
        if seedfile:
            extgenlist = tp.list_of_names_in_line(open(seedfile))
            # Index of genes belonging to the "universe" which also have var .=lvar and med >=lmedS
            #print 'Genes in matrix expMat'
            #print expMat.genInd
            
            dictmgl2, genenotfound = sd.get_gen_index(expMat,extgenlist)
            #mgenelist, dictmgl =  tl.get_variable_genes(expMat,finalgenes)
            #print genenotfound
            
            #print 'Genes in present in the universe list'
            #print dictmgl2 
            expMat = tl.get_filtered_expMat(expMat,dictmgl2,len(samplelist))
        
        #expMat.normalization()
        expMat.log2_norm()
        #print mgenelist
        #print 'Genes not found in the variable genes'        
        #print genenotfound
        #print 'Number of genes ' + str(genenotfound)
        
        return expMat 




def get_expMat(pipe, samplelist, genelist, norm=False):
    """
    Returns a expMat, an expression Array object (common classes) object
    input: pipe= pipeline object, samplelist = a list of samples, genelist=list of genes
    norm = bool indicates if normalization of expression data or or just log2(expression)
    """
    
    pipe.mq.create_connector()
    expMat = pipe.mq.get_exp_mat_2(samplelist, genelist, pipe.tableName)
    # Normalize the expression values data
    if norm:
        expMat.normalization()
    else:
        expMat.log2_norm()
        
    return expMat


def write_expmat_sdv(pipe, expMat, genelistind, outfile1, outfile2, spclass0, spclass1=[], spcltest=[]):
    """
    writes a expression matrix. The matrix is ordered according to the sample clasess: spclass0, spclass1, spcltest
    expMat= an expression Array object (common classes), genelistind =  the index in expMat of the wanted 
    genes.  string outfile1= writting file for the gene list, outfile2=writting file for the expression matrix
    """
    pipe.mq.create_connector()
    pipe.mq.write_variable_genes_sdv(expMat, genelistind, outfile1, outfile2, spclass0, spclass1, spcltest)



##########

if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--annotFile", dest="annotFile",
                            help="annotation file for all files to use")
    optionparser.add_option("-d", action="store_true", dest="dbtrue",
                            help="create a database. Bool")
    
    optionparser.add_option("-n", action="store_true", dest="norm",
                            help="normalize log2 and scale expression vals. Default only log2. Default universe. Bool")

    
    (options, args) = optionparser.parse_args()
       
    # start the pipeline. 
    # open the database
    pipe = pipeline("localhost", "oabalbin", "oscar", "signatures")
    expr_folder, seed_file, db_file, samples, sampleclass, outdirs = pipe.parse_input_annotfiles(open(options.annotFile))    
    #read folder for expression profiles
    exp_files = pipe.read_files_folder(expr_folder['expdir'],'.txt')
    
    ####### output files
    
    outfile1 = outdirs['outdir']+'svd/mat_exp_SiRNA_NCI60_5.txt'   # bfrm parameters file
    outfile2 = outdirs['outdir']+'svd/gen_exp_SiRNA_NCI60_5.txt'   # bfrm parameters file

    ######### Get raw data  ############
    # Get the database files
    tp = pm.parser()
    pipe.mt.create_connector()
    pipe.mq.create_connector()
    gtp = gpm.geneparser()
    
    if options.dbtrue:
        pipe.initialize_database(exp_files,tp, options.dbtrue)

    
    ########## Make queries ############
    
    # list of samples:
    samplelist = tp.list_of_names_in_line(open(samples['samples']))
    genelist = tp.list_of_names(open(seed_file['univfile']))    
    
    expMat = get_expMat(pipe, samplelist, genelist, options.norm)

        
    ########### SDV
    # write input for the sdv program
    forallProbes=range((expMat.expVal.shape[0]))
    write_expmat_sdv(pipe, expMat, forallProbes, open(outfile1,'w'), open(outfile2,'w'), sampleclass['spclass0'], sampleclass['spclass1'],sampleclass['spcltest'])

    sys.exit(0)
